from kafka import KafkaProducer, KafkaConsumer, TopicPartition
from PyQt5.QtCore import QThread, pyqtSignal, QTimer


class KafkaProducerThread(QThread):
    message_sent = pyqtSignal(str)  # 定义信号

    def __init__(self, parent=None):
        """
        初始化
        @param parent:
        """
        super().__init__(parent)
        self.producer = None
        self._run = True
        self._brokers = "localhost:9092"
        self._queue = None
        self._topic = None
        self._msg = "hello kafka"

    def set_brokers(self, brokers):
        """
        设置kafka服务器地址
        @param brokers:
        @return:
        """
        self._brokers = brokers

    def set_queue(self, queue):
        """
        设置队列
        @param queue:
        @return:
        """
        self._queue = queue

    def set_topic(self, topic):
        """
        设置主题
        @param topic:
        @return:
        """
        self._topic = topic

    def set_msg(self, msg):
        """
        设置消息
        @param msg:
        @return:
        """
        self._msg = msg

    def stop(self):
        """
        停止线程
        @return:
        """
        self._run = False

    def close(self):
        """
        关闭生产者
        @return:
        """
        self.producer.close()

    def finished(self):
        """
        线程结束时调用
        @return:
        """
        self.producer.close()

    def run(self):
        """
        重写run方法
        @return:
        """
        if self._queue:
            self.producer = KafkaProducer(bootstrap_servers=self._brokers)
            self.producer.send(self._queue, self._msg.encode())
        elif self._topic:
            self.producer = KafkaProducer(bootstrap_servers=self._brokers)
            self.producer.send(self._topic, self._msg.encode())
        self.message_sent.emit(self._msg)


class KafkaConsumerThread(QThread):
    message_received = pyqtSignal(str)  # 自定义信号，执行run()函数时，从相关线程发射此信号

    def __init__(self, parent=None):
        """
        初始化
        @param parent:
        """
        super().__init__(parent)
        self.consumer = None
        self._run = True
        self._brokers = "localhost:9092"
        self._queue = None
        self._topic = None
        self._partition = None

    def set_brokers(self, brokers):
        """
        设置kafka服务器地址
        @param brokers:
        @return:
        """
        self._brokers = brokers

    def set_queue(self, queue):
        """
        设置队列
        @param queue:
        @return:
        """
        self._queue = queue

    def set_topic(self, topic):
        """
        设置主题
        @param topic:
        @return:
        """
        self._topic = topic

    def set_partition(self, partition):
        """
        设置分区
        @param partition:
        @return:
        """
        self._partition = partition

    def close(self):
        """
        关闭消费者
        @return:
        """
        self.consumer.close()

    def stop(self):
        """
        停止线程
        @return:
        """
        self._run = False

    def run(self):
        """
        重写run方法
        @return:
        """
        if self._queue:  # 队列
            self.consumer = KafkaConsumer(
                bootstrap_servers=self._brokers,
                group_id="default-group",
                auto_offset_reset="earliest",
                enable_auto_commit=True,
                auto_commit_interval_ms=1000,
            )
            self.consumer.subscribe(self._queue)
        elif self._topic:  # 主题
            self.consumer = KafkaConsumer(
                bootstrap_servers=self._brokers,
                group_id="default-group",
                enable_auto_commit=True,
                auto_commit_interval_ms=1000,
            )
            if self._partition:
                topic_partition = TopicPartition(self._topic, self._partition)
                self.consumer.assign([topic_partition])
            else:
                self.consumer.subscribe([self._topic])
        for message in self.consumer:
            self.message_received.emit(message.value.decode())


if __name__ == "__main__":
    pass
